feat(algorithms): add AR PPO with GAE and value loss - #2
Open
yhl48 wants to merge 12 commits into
Open
Conversation
Wire per-token critic values through Qwen3 replay and compute GAE advantages on RolloutTrack for the PPO critic path (issue Tencent-Hunyuan#86, part 2/3).
Keep transformers imports inside from_config / _packed_replay_supported as in upstream; only UniRL imports (ValueHead, ARReplayOutput) stay top-level.
Call _require_value_head_for_replay at the start of _replay_aware_forward so return_values=True fails fast before the transformer forward, not only from Qwen3ARStage.replay().
Extend ReplayResult with optional per-token values for PPO/GAE and replace ARReplayOutput in Qwen3 replay with the shared type.
Apply ruff-format changes and remove unused cu in _finalize_replay_output.
Derive last-token positions from packed offsets (end - 1) instead of passing redundant lengths, with a guard for empty samples.
yhl48
force-pushed
the
feat/ar-value-head
branch
from
July 26, 2026 16:44
87c5630 to
6728292
Compare
yhl48
force-pushed
the
feat/ar-value-head
branch
from
July 26, 2026 20:13
6728292 to
bf54bf5
Compare
leviking98z-rgb
force-pushed
the
feat/ar-value-head
branch
from
August 1, 2026 17:03
bf54bf5 to
c6aae06
Compare
leviking98z-rgb
force-pushed
the
feat/ar-ppo-gae
branch
from
August 1, 2026 17:03
c6c7681 to
c8f1afe
Compare
Co-authored-by: leviking98z-rgb <leviking98z@gmail.com>
Wire PPO.prepare_rollout_track for worker-side GAE, policy+value losses, advantage_mode=gae in AR trainers, and a Qwen3 DAPO recipe (issue Tencent-Hunyuan#86, 3/3).
leviking98z-rgb
force-pushed
the
feat/ar-ppo-gae
branch
2 times, most recently
from
August 1, 2026 18:43
7acd52d to
cda2582
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds the PPO algorithm path for autoregressive training (issue Tencent-Hunyuan#86, part 3/3), stacked on
feat/ar-value-head(Tencent-Hunyuan#256 upstream).PPO/PPOConfigwith worker-sideprepare_rollout_track(no-grad critic replay + GAE) and train-step policy clip + clipped value loss._ppo_clipped_value_lossinalgorithms/base.py.TrainStackcalls optionalalgorithm.prepare_rollout_track()before micro-batch prep.advantage_mode: grpo | gaeon AR trainers.examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml.Related Issue
Part of Tencent-Hunyuan#86 (PPO value critic + GAE for AR models).
Test Plan
python3 -m ruff check unirl/algorithms/ppo.py unirl/algorithms/base.py unirl/algorithms/__init__.py unirl/train/stack/base.py unirl/train_ar.py unirl/trainer/ar.py unirl/trainer/async_ar.py tests/algorithms/test_ppo_value_loss.pypytest tests/algorithms/test_ppo_value_loss.py(Python 3.12+ project env)Compatibility / Risk
advantage_mode=grpounchanged.use_value_head=Trueandadvantage_mode: gae.value_head.*for separate SGLang engine.Reviewer Notes
Tencent-Hunyuan/UniRLwith basefeat/ar-value-headonce feat(models): add AR value head and GAE wiring Tencent-Hunyuan/UniRL#256 branch is available there, or merge feat(models): add AR value head and GAE wiring Tencent-Hunyuan/UniRL#256 first then rebase ontomain.old_logp,old_values, GAE advantages/returns; moving per update:new_logp,new_values.Checklist